import pendulum
from datetime import timedelta
from utils.operators.rich_sql_sensor import RichSqlSensor
from airflow.exceptions import AirflowSkipException

# 量本利 RichSql 非8号 15号 跳过运行
cst = pendulum.timezone('Asia/Shanghai')

class VolumeRichSqlSensor(RichSqlSensor):

    def pre_execute(self, context):
        day = cst.convert(context['ti'].execution_date) + timedelta(days=1)

        schedule_date = ['05','15','22']

        if day.strftime('%d') not in schedule_date:
            print(f'{day.strftime("%d")} not in {schedule_date}, should skip')
            raise AirflowSkipException()
        else:
            print(f'{day.strftime("%d")} in {schedule_date}, run now')
            super().pre_execute(context)